import asyncio
import json
import re
import time

from llm_model.llm_model import LLMModel

from tools.config import extract_text_from_file
from app.db.es_db import ElasticsearchClient
from app.db.noe4j_db import Neo4jClient


class TextSplit:
    def __init__(self):
        self.llm = LLMModel("kimi")
        self.extract_text_from_file = extract_text_from_file
        self.es = ElasticsearchClient(hosts=["http://localhost:9200"], username="elastic", password="elastic")
        self.noe4j = Neo4jClient(uri="bolt://localhost:7687", user="neo4j", password="neo4j_test")
        self.role_content = self.get_role_content()

    @staticmethod
    def get_role_content():
        with open('llm_model/prompt/dian_li.txt', 'r', encoding='utf-8') as f:
            data = f.read()
        return data

    def get_text(self, file_path: str) -> list:
        data = self.extract_text_from_file(file_path)
        data = data.split('。')
        return data

    @staticmethod
    def get_json(data):
        data = re.findall('\{.+}', data, re.S)
        if data:
            return data[0].replace('\'', '"')

    def store_text_in_es(self, index_name, text):
        """
        将文本存储到Elasticsearch中
        :param index_name: 索引名称
        :param text: 要存储的文本
        """
        # 创建索引（如果不存在）
        self.es.create_index(index_name)

        # 插入文档
        document = {
            "text": text
        }
        response = self.es.insert_document(index_name, document)
        print(f"文档存储成功: {response}")

    def store_triple_in_neo4j(self, triple):
        """
        将SPO形式的三元组存储到Neo4j中
        :param neo4j_client: Neo4jClient实例
        :param triple: 三元组（字典格式，包含主体、谓语和对象）
        """
        subject = triple['subject']
        predicate = triple['predicate']
        object = triple['object']
        query = (
            f"MATCH (a:{subject['label']} {{name: $subject_name}})-[r:{predicate}]->(b:{object['label']} {{name: $object_name}}) "
            f"RETURN count(r) as relationship_count"
        )
        result = self.noe4j.run_query(query, {"subject_name": subject["properties"]["name"],
                                              "object_name": object["properties"]["name"]})
        relationship_count = result[0]['relationship_count']

        if relationship_count > 0:
            print("相同的关系已存在，跳过存储。")
            return

        # 创建主体节点
        subject_node = self.noe4j.create_node(subject["label"], subject["properties"])
        print(f"主体节点创建成功: {subject_node}")

        # 创建对象节点
        object_node = self.noe4j.create_node(object["label"], object["properties"])
        print(f"对象节点创建成功: {object_node}")

        # 创建关系
        rel = self.noe4j.create_relationship(
            node1_label=subject["label"],
            node1_key=subject["properties"]["name"],
            node2_label=object["label"],
            node2_key=object["properties"]["name"],
            rel_type=predicate,
            properties=None  # 谓语通常不包含属性，但可以根据需要添加
        )
        print(f"关系创建成功: {rel}")

    async def get_three_tuple(self, data_list: list):
        for data in data_list:
            self.store_text_in_es("dian_li", data)
            time.sleep(30)
            result = await self.llm.chat(query=data,
                                         role_content=self.role_content)
            try:
                result = self.get_json(result)
                self.store_triple_in_neo4j(json.loads(result))
                print(f'存储关系{result}')
            except Exception as e:
                print(e)

    async def start(self, file_path: str):
        data_list = self.get_text(file_path=file_path)
        await self.get_three_tuple(data_list)
        return True


if __name__ == '__main__':
    asyncio.run(TextSplit().start('data/dianli/电力知识.docx'))
